import os
import torch
import numpy as np
from scipy.spatial import Delaunay
import json

class_to_name = {
    0: 'car',
    1: 'motorcycle',
    2: 'bus',
    3: 'truck',
    4: 'bicycle',
    5: 'pedestrians',
    6: 'Engineering_vehicles',
    7: 'Traffic_cone',
    8: 'trailer',
    9: 'police_car',
    10: 'A_push-pull_object',
    11: 'stroller',
    12: 'Portable_personal_motor_vehicle',
    13: 'Temporary_traffic_obstruction',
    14: 'wheelchair',
    15: 'ambulance',
    16: 'animal',
    17: 'other',
    18: 'gt_label',
    19: 'pr_label'
}
for i in class_to_name.keys():
    class_to_name[i] = class_to_name[i].lower()
name_to_class = dict(zip(class_to_name.values(),class_to_name.keys()))
class object3d():
    def __init__(self, label):
        self.type = label[1]
        self.visible_range = [int(i) for i in label[2].split('—')]
        self.status = label[3]
        self.x = float(label[4])
        self.y = float(label[5])
        self.z = float(label[6])
        self.l = float(label[7])
        self.w = float(label[8])
        self.h = float(label[9])
        self.rz = float(label[10])

def check_numpy_to_torch(x):
    if isinstance(x, np.ndarray):
        return torch.from_numpy(x).float(), True
    return x, False
def rotate_points_along_z(points, angle):
    """
    Args:
        points: (B, N, 3 + C)
        angle: (B), angle along z-axis, angle increases x ==> y
    Returns:

    """
    points, is_numpy = check_numpy_to_torch(points)
    angle, _ = check_numpy_to_torch(angle)

    cosa = torch.cos(angle)
    sina = torch.sin(angle)
    zeros = angle.new_zeros(points.shape[0])
    ones = angle.new_ones(points.shape[0])
    rot_matrix = torch.stack((
        cosa,  sina, zeros,
        -sina, cosa, zeros,
        zeros, zeros, ones
    ), dim=1).view(-1, 3, 3).float()
    points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
    return points_rot.numpy() if is_numpy else points_rot
def in_hull(p, hull):
    """
    :param p: (N, K) test points
    :param hull: (M, K) M corners of a box
    :return (N) bool
    """
    try:
        if not isinstance(hull, Delaunay):
            hull = Delaunay(hull)
        flag = hull.find_simplex(p) >= 0
    except scipy.spatial.qhull.QhullError:
        print('Warning: not a hull %s' % str(hull))
        flag = np.zeros(p.shape[0], dtype=np.bool)

    return flag


def boxes_to_corners_3d(boxes3d):
    """
        7 -------- 4
       /|         /|
      6 -------- 5 .
      | |        | |
      . 3 -------- 0
      |/         |/
      2 -------- 1
    Args:
        boxes3d:  (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center

    Returns:
    """
    boxes3d, is_numpy = check_numpy_to_torch(boxes3d)

    template = boxes3d.new_tensor((
        [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1],
        [1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1],
    )) / 2

    corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
    corners3d = rotate_points_along_z(corners3d.view(-1, 8, 3), boxes3d[:, 6]).view(-1, 8, 3)
    corners3d += boxes3d[:, None, 0:3]

    return corners3d.numpy() if is_numpy else corners3d



def get_box_label_txt(path_label):
    with open(path_label) as f:
        lines = f.readlines()
    obj_list = [object3d(i.split()) for i in lines]
    annotations = {}
    annotations['name'] = np.array([obj.type for obj in obj_list])
    annotations['dimensions'] = np.array([[obj.l, obj.w, obj.h] for obj in obj_list])
    annotations['location'] = np.array([[obj.x, obj.y, obj.z] for obj in obj_list])
    annotations['rotation_z'] = np.array([obj.rz for obj in obj_list])
    annotations['status'] = [obj.status for obj in obj_list]
    annotations['visible_range'] = [obj.visible_range for obj in obj_list]

    num_gt = len(annotations['name'])
    loc = annotations['location']
    dims = annotations['dimensions']
    rots = annotations['rotation_z'].reshape(-1, 1)
    gt_boxes_lidar = np.concatenate([loc, dims, rots], axis=1)
    annotations['gt_boxes_lidar'] = gt_boxes_lidar
    name = [i.lower() for i in annotations['name']]
    return name, gt_boxes_lidar

def get_box_label_json(path_label):
    # with open(path_label, 'r') as f:
    #     data = json.load(f)
    # name = [i['obj_type'].lower() for i in data]
    # xyz = np.array(
    #     [[i['psr']['position']['x'], i['psr']['position']['y'], i['psr']['position']['z']] for i in data])
    # lwh = np.array([[i['psr']['scale']['x'], i['psr']['scale']['y'], i['psr']['scale']['z']] for i in data])
    # r = np.array([[i['psr']['rotation']['z']] for i in data])
    # box = np.concatenate([xyz, lwh, r], axis=-1)

    with open(path_label, 'r') as f:
        data_label_pr = json.load(f)
    # print(data_label_pr)
    name = data_label_pr['annos']['name']
    box = data_label_pr['annos']['gt_boxes_3d']
    box = np.array(box)
    # score = data_label_pr['annos']['confidence']
    # velocity_pr = data_label_pr['annos']['velocity']
    # acceleration_pr = data_label_pr['annos']['acceleration']
    return name,box

def get_box_label(path_label):
    if path_label[-3:]=='txt':
        name, box = get_box_label_txt(path_label)
    else:
        name, box = get_box_label_json(path_label)
    return name,box

def get_lidar(path_lidar):
    points = np.fromfile(str(path_lidar),dtype=np.float32).reshape(-1,4)
    return points

def Get_seglabel(name,gt_boxes_lidar,points):
    corners_lidar = boxes_to_corners_3d(gt_boxes_lidar)
    label = np.array([name_to_class[i] + 1 for i in name])
    # mask = np.array([True for i in name if i in name_to_class.keys() else False])
    # name = [i for i in name if i in name_to_class.keys()]
    # types = np.array([name_to_class[i] for i in name if i in name_to_class.keys()]).reshape(-1,1)
    gt_boxes_lidar = gt_boxes_lidar
    seg_types = np.zeros((points.shape[0],), dtype=np.float32)
    num_gt = len(name)
    for k in range(num_gt):
        flag = in_hull(points[:, 0:3], corners_lidar[k])
        seg_types += flag * (name_to_class[name[k]] + 1)
    return label,corners_lidar, seg_types

def main(data_file):
    pts_file = data_file + '/data_bin'
    labels_file = data_file + '/label'
    labels_seg_out = data_file + '/label_bin'
    pkg_list = os.listdir(pts_file)
    for pkg_list_ in pkg_list:
        pts_list = os.listdir(pts_file + '/' + pkg_list_ + '/lidar')
        for pts in pts_list:
            pt_file = pts_file +'/' + pkg_list_ + '/lidar/' + pts
            label_file = labels_file +'/' + pkg_list_ + '/lidar/all/' + pts[:-8] + '.txt'
            label_out_path = labels_seg_out +'/' + pkg_list_ + '/lidar/'
            if not os.path.exists(label_out_path):
                os.makedirs(label_out_path)
            label_out = label_out_path + pts[:-8] + '.bin'

            pts = get_lidar(pt_file)
            name,box = get_box_label(label_file)
            corners_lidar,seg_types = Get_seglabel(name,box,pts)
            seg_types = seg_types.astype(np.uint8)
            seg_types.tofile(label_out)
            print(label_out)
if __name__ == '__main__':
    pass
    # data_file = '/media/zhl/00030FC90009D3A0/OpenPCDet/data/livox_dataset_v1.0/'
    # main(data_file)
